#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
网页全屏截图工具
支持对包含JavaScript动态滚动条的网页进行完整截图
"""

import argparse
import time
import os
import sys
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.chrome.service import Service as ChromeService
from selenium.common.exceptions import WebDriverException
from PIL import Image
import io

class WebScreenshot:
    def __init__(self):
        """初始化WebScreenshot类"""
        try:
            # 设置Chrome选项
            self.chrome_options = webdriver.ChromeOptions()
            self.chrome_options.add_argument('--headless=new')  # 使用新的无界面模式
            self.chrome_options.add_argument('--start-maximized')
            self.chrome_options.add_argument('--disable-gpu')
            self.chrome_options.add_argument('--no-sandbox')
            self.chrome_options.add_argument('--disable-dev-shm-usage')
            self.chrome_options.add_argument('--ignore-certificate-errors')  # 忽略证书错误
            self.chrome_options.add_argument('--ignore-ssl-errors')  # 忽略SSL错误
            
            print("正在初始化Chrome浏览器...")
            
            # 直接使用Chrome浏览器
            self.driver = webdriver.Chrome(
                options=self.chrome_options
            )
            
            print("Chrome浏览器初始化成功")
            
        except WebDriverException as e:
            print(f"Chrome浏览器初始化失败: {str(e)}")
            print("请确保：")
            print("1. 已安装最新版本的Chrome浏览器")
            print("2. 系统环境变量PATH中包含Chrome浏览器的路径")
            print("3. 如果问题仍然存在，请尝试重新安装Chrome浏览器")
            sys.exit(1)
        except Exception as e:
            print(f"初始化失败: {str(e)}")
            print("详细错误信息:", sys.exc_info())
            sys.exit(1)

    def get_full_screenshot(self, url: str, output_path: str, format: str = 'PNG', quality: int = 90) -> None:
        """
        获取网页完整截图
        
        Args:
            url: 网页URL
            output_path: 输出图片路径
            format: 图片格式（PNG/JPEG）
            quality: 图片质量（仅对JPEG有效）
        """
        try:
            print(f"正在访问网页: {url}")
            self.driver.get(url)
            
            print("等待页面加载...")
            WebDriverWait(self.driver, 20).until(
                EC.presence_of_element_located((By.TAG_NAME, "body"))
            )
            
            # 等待页面完全加载
            time.sleep(5)  # 增加等待时间，确保页面完全加载
            
            # 获取页面实际高度
            print("计算页面尺寸...")
            total_height = self.driver.execute_script("""
                return Math.max(
                    document.body.scrollHeight,
                    document.documentElement.scrollHeight,
                    document.body.offsetHeight,
                    document.documentElement.offsetHeight,
                    document.body.clientHeight,
                    document.documentElement.clientHeight
                );
            """)
            
            # 设置窗口大小
            width = self.driver.execute_script("return document.documentElement.scrollWidth")
            print(f"设置窗口大小: {width}x{total_height}")
            self.driver.set_window_size(width, total_height)
            
            # 等待页面渲染
            print("等待页面渲染完成...")
            time.sleep(2)
            
            # 滚动到页面底部确保所有内容都加载
            self.driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
            time.sleep(2)  # 等待滚动后的内容加载
            self.driver.execute_script("window.scrollTo(0, 0);")  # 回到顶部
            time.sleep(1)
            
            # 获取截图
            print("正在截图...")
            screenshot = self.driver.get_screenshot_as_png()
            
            # 使用Pillow处理图片
            image = Image.open(io.BytesIO(screenshot))
            
            # 保存图片
            print(f"正在保存图片到: {output_path}")
            image.save(output_path, format=format, quality=quality)
            print(f"截图已成功保存到: {output_path}")
            
        except WebDriverException as e:
            print(f"浏览器操作错误: {str(e)}")
            print("请检查网络连接和URL是否正确")
        except Exception as e:
            print(f"截图过程中发生错误: {str(e)}")
            print("详细错误信息:", sys.exc_info())
        finally:
            self.close()
    
    def close(self):
        """关闭浏览器"""
        try:
            if hasattr(self, 'driver') and self.driver:
                self.driver.quit()
                print("浏览器已关闭")
        except Exception as e:
            print(f"关闭浏览器时发生错误: {str(e)}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='网页全屏截图工具')
    parser.add_argument('--url', required=True, help='要截图的网页URL')
    parser.add_argument('--output', default='screenshot.png', help='输出图片路径')
    parser.add_argument('--format', default='PNG', choices=['PNG', 'JPEG'], help='输出图片格式')
    parser.add_argument('--quality', type=int, default=90, help='图片质量(仅对JPEG格式有效)')
    
    args = parser.parse_args()
    
    try:
        screenshot = WebScreenshot()
        screenshot.get_full_screenshot(args.url, args.output, args.format, args.quality)
    except KeyboardInterrupt:
        print("\n程序被用户中断")
        sys.exit(0)
    except Exception as e:
        print(f"程序执行出错: {str(e)}")
        print("详细错误信息:", sys.exc_info())
        sys.exit(1)

if __name__ == '__main__':
    main() 